package WSTClient;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.ByteBuffer;
import java.util.concurrent.atomic.AtomicBoolean;

import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.crypto.asymmetric.KeyType;
import cn.hutool.crypto.asymmetric.RSA;
import cn.hutool.crypto.symmetric.RC4;

public class WSClient {

	public static void main(String[] args) {
		
		String pub = FileUtil.readString("pub", CharsetUtil.CHARSET_GBK);
		RSA rsa = new RSA(null,pub);
		
		String k = RandomUtil.randomString(16);
		RC4 rc = new RC4(k);
		
		try (
				ServerSocket server = new ServerSocket(Integer.valueOf(args[0]));
			){
			
			while(true) {
				
				Socket skt = server.accept();
				
				InputStream in = skt.getInputStream();
				OutputStream out = skt.getOutputStream();
				
				new Thread( () -> {
					buildLink(args[1],args[2],k,skt,in, out, rc, rsa);
				} ).start();
			}
			
		} catch (IOException e) {
			e.printStackTrace();
		}
		
	}

	public static void buildLink(String tar,String ws,String k,Socket skt,InputStream input, OutputStream out,RC4 rc,RSA rsa) {
		
		AtomicBoolean wconnect = new AtomicBoolean(true);
		AtomicBoolean wsc = new AtomicBoolean(true);
		
		WebSocketClient client = null;
		try {
			client = new WebSocketClient(new URI(ws)) {

				@Override
				public void onOpen(ServerHandshake handshakedata) {
					
				}

				@Override
				public void onMessage(String message) {
					//接收到回复后,标记步骤完成
					wconnect.set(false);
				}
				
				@Override
				public void onMessage(ByteBuffer bytes) {
					byte[] b = rc.crypt(bytes.array());
					try {
						out.write(b);
					} catch (IOException e) {
						e.printStackTrace();
					}
				}

				@Override
				public void onError(Exception ex) {
					wsc.set(false);
				}

				@Override
				public void onClose(int code, String reason, boolean remote) {
					wsc.set(false);
				}
			};
		} catch (URISyntaxException e2) {
			e2.printStackTrace();
		}
		
		try{
			
			client.connectBlocking();
			
			//建立连接后，发送信息
			client.send(rsa.encryptHex(k, KeyType.PublicKey));
			
			while(wconnect.get()) {
				Thread.sleep(200);
			}
			
			//收到回复后，发送连接目标
			client.send(rc.encryptHex(tar));
			wconnect.set(true);
			
			while(wconnect.get()) {
				Thread.sleep(200);
			}
			
			wconnect.set(true);
			
			//收到连接建立后,建立本地连接
			while(wsc.get()) {
				
				byte[] b = new byte[1024];
				int len = 0;
				
				while((len = input.read(b) ) > 0) {
					byte[] send = new byte[len];
					System.arraycopy(b, 0, send, 0, len);
					client.send(rc.crypt(send));
				}
				
			}
			
		} catch (Exception e) {
			e.printStackTrace();
		}finally {
			wsc.set(false);
			try {
				client.close();
				skt.close();
			} catch (IOException e1) {
				e1.printStackTrace();
			}
			
		}
	}

}
